import pandas as pd
import datetime
import matplotlib.pyplot as plt
import seaborn as sns
from utils.log import Logger
from sklearn.ensemble import RandomForestClassifier
from sklearn.feature_selection import RFECV
from sklearn.model_selection import train_test_split, StratifiedKFold
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
import warnings

warnings.filterwarnings("ignore")
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


class FeatureEngineering(object):
    def __init__(self, path):
        """初始化"""
        # 配置日志
        train_log_name = "FeatureEngineering_" + datetime.datetime.now().strftime("%Y%m%d%H%M%S")
        self.logger = Logger('../', train_log_name).get_logger()
        # 获取源数据
        self.data_source = pd.read_csv(path)
        # 初始化结果存储
        self.precessed_data = None  # 存储预处理处理后的数据
        self.x = None  # 存储特征
        self.y = None  # 存储标签
        self.selected_features = None  # 存储选择的特征
        self.feature_importance = None  # 存储特征重要性

        self.logger.info("特征工程类初始化完成".center(50, '='))

    def preprocess(self):
        """
        特征预处理
        1.删除无关列数据
        2.类别列特征one-hot编码
        :return: 特征x和标签y
        """
        self.logger.info("开始特征预处理".center(50, '='))

        data = self.data_source.copy()

        # 删除 固定值列 和 明显无关列 数据
        data.drop(['EmployeeNumber', 'StandardHours', 'Over18'], axis=1, inplace=True)

        # 类别列进行one-hot编码
        categorical_cols = data.select_dtypes(include=['object']).columns.tolist()
        data = pd.get_dummies(data, columns=categorical_cols, drop_first=True)

        # 特征交叉工程
        self.logger.info("开始特征交叉工程")

        # 年龄收入比：衡量单位年龄的收入水平，可能反映员工的薪资增长情况
        # 值越高表示同龄人中收入越高，可能影响员工离职倾向
        if 'Age' in data.columns and 'MonthlyIncome' in data.columns:
            data['Age_Income_Ratio'] = data['Age'] / (data['MonthlyIncome'] / 1000 + 1e-8)  # 添加小值避免除零
            self.logger.info("添加特征：Age_Income_Ratio")

        # 工作年限收入比：衡量单位工作年限的收入水平，反映工作经验的回报率
        # 可能影响员工对当前工作的满意度
        if 'TotalWorkingYears' in data.columns and 'MonthlyIncome' in data.columns:
            data['Years_Income_Ratio'] = data['TotalWorkingYears'] / (data['MonthlyIncome'] / 1000 + 1e-8)
            self.logger.info("添加特征：Years_Income_Ratio")

        # 晋升间隔：计算距离上次晋升的间隔时间
        # 间隔越长可能表示晋升机会少，影响员工留存
        if 'YearsAtCompany' in data.columns and 'YearsSinceLastPromotion' in data.columns:
            data['Promotion_Gap'] = data['YearsAtCompany'] - data['YearsSinceLastPromotion']
            self.logger.info("添加特征：Promotion_Gap")

        # 角色任期比：当前角色任期占在公司总时间的比例
        # 比值低可能表示频繁变动角色或新员工，比值高可能表示长期固定在一个角色
        if 'YearsInCurrentRole' in data.columns and 'YearsAtCompany' in data.columns:
            data['Role_Tenure_Ratio'] = data['YearsInCurrentRole'] / (data['YearsAtCompany'] + 1e-8)
            self.logger.info("添加特征：Role_Tenure_Ratio")

        # 经验效率比：单位年龄的工作经验，衡量经验积累的效率
        # 可反映员工的职业成熟度和早期职业发展情况
        if 'TotalWorkingYears' in data.columns and 'Age' in data.columns:
            data['Experience_Age_Ratio'] = data['TotalWorkingYears'] / (data['Age'] + 1e-8)
            self.logger.info("添加特征：Experience_Age_Ratio")

        # 收入增长潜力：单位工作年限的收入，衡量收入随经验增长的潜力
        # 反映员工在当前公司的收入增长情况
        if 'MonthlyIncome' in data.columns and 'TotalWorkingYears' in data.columns:
            data['Income_Growth_Potential'] = data['MonthlyIncome'] / (data['TotalWorkingYears'] + 1e-8)
            self.logger.info("添加特征：Income_Growth_Potential")

        self.logger.info(f"特征交叉后数据形状：{data.shape}")

        # 分离特征和标签
        self.x = data.drop('Attrition', axis=1)
        self.y = data['Attrition']

        self.precessed_data = data
        self.logger.info(f"预处理后数据形状：{data.shape}")
        self.logger.info(f"预处理后特征数量：{self.x.shape[1]}")
        self.logger.info("特征预处理结束".center(50, '='))

        return self.x, self.y

    def feature_importance_rf(self):
        """
        使用随机森林计算特征重要性
        :return:
        """
        self.logger.info("使用随机森林计算特征重要性".center(50, '='))

        # 训练随机森林模型
        rf = RandomForestClassifier(n_estimators=100, random_state=11, n_jobs=-1)
        rf.fit(self.x, self.y)

        # 计算特征重要性
        importance = pd.Series(rf.feature_importances_, index=self.x.columns).sort_values(ascending=False)
        self.feature_importance = importance

        # 转换为DataFrame
        importance_df = importance.reset_index()
        importance_df.columns = ['Feature', 'Importance']

        self.logger.info(f"特征重要性排序: \n{importance_df}")

        # 可视化
        plt.figure(figsize=(20, 10))
        # 根据重要性值创建颜色渐变
        colors = plt.cm.Reds(importance.values / importance.values.max())
        plt.barh(importance.index, importance.values, color=colors)
        plt.title("随机森林特征重要性")
        plt.xlabel("重要性")
        plt.gca().invert_yaxis()  # 保持重要性从高到低的顺序
        plt.tight_layout()
        plt.grid(axis='x', alpha=0.7)
        # 保存图像
        plt.savefig('../data/fig/rf_feature_importance.png', dpi=300)
        plt.show()

        self.logger.info("重要性计算完成".center(50, '='))

        return importance_df

    def select_features_rfe(self, n_estimators=100, random_state=11, cv=5, scoring='f1', step=1):
        """使用递归特征消除(RFECV)进行特征选择
        Args:
            n_estimators: 随机森林的树数量
            random_state: 随机种子
            cv: 交叉验证折数
            scoring: 评分标准
            step: 每次移除的特征数量
        Returns:
            选择的特征列表
        """
        self.logger.info("开始 RFECV+随机森林 特征选择".center(50, "="))

        # 构建Pipeline（防止数据泄露）
        self.logger.info("构建Pipeline，包含StandardScaler和RFECV")
        pipeline = Pipeline([
            ('scaler', StandardScaler()),  # 标准化
            ('rfecv', RFECV(
                estimator=RandomForestClassifier(n_estimators=n_estimators, random_state=random_state, n_jobs=-1),
                step=step,
                cv=StratifiedKFold(cv, shuffle=True, random_state=random_state),
                scoring=scoring,
                n_jobs=-1,
                min_features_to_select=18
            ))
        ])

        # 拟合模型
        pipeline.fit(self.x, self.y)
        selector = pipeline.named_steps['rfecv']

        # 获取选择的特征
        selected_features = self.x.columns[selector.support_]
        self.selected_features = selected_features.tolist()

        # 特征排名
        features_rank = pd.DataFrame({
            'feature': self.x.columns,
            'rank': selector.ranking_
        }).sort_values('rank')

        # 记录结果
        self.logger.info(f"最佳特征数量: {selector.n_features_}")
        self.logger.info(f"被选中的特征: {', '.join(self.selected_features)}")
        self.logger.info(f"特征重要性排名: \n{features_rank}")

        # 可视化
        self.logger.info("绘制特征选择结果图")
        plt.figure(figsize=(10, 6))
        cv_scores = selector.cv_results_['mean_test_score']
        plt.plot(range(1, len(cv_scores) + 1), cv_scores, marker='o')
        plt.xlabel("选择的特征数量")
        plt.ylabel(f"{cv}折交叉验证的{scoring}分数")
        plt.title("递归特征消除(RFECV)特征选择")
        plt.grid()
        plt.axvline(selector.n_features_, color='r', linestyle='--', label=f'最优特征数 = {selector.n_features_}')
        plt.legend()
        plt.tight_layout()
        # 保存图像
        plt.savefig('../data/fig/RFESCV_best_features.png', dpi=300)
        plt.show()

        self.logger.info("结束 RFECV+随机森林 特征选择".center(50, "="))
        return self.selected_features

    def features_corr(self):
        """
        计算进行RFECV之后特征之间的皮尔逊系数
        去除相关性过高的特征之一
        :return:
        """
        self.logger.info("开始计算特征间的相关性".center(50, "="))

        # 检查是否已经进行了特征选择
        if self.selected_features is None:
            self.logger.warning("尚未进行特征选择，使用所有特征进行相关性分析")
            features_data = self.x
        else:
            self.logger.info(f"使用RFECV选择的{len(self.selected_features)}个特征进行相关性分析")
            features_data = self.x[self.selected_features]

        # 计算相关性矩阵
        corr_matrix = features_data.corr().abs()

        # 创建下三角矩阵掩码
        """
        下三角区域：保留原始的相关性系数值
        上三角区域（包括对角线）：值被替换为NaN
        """
        lower_triangle_mask = pd.DataFrame(
            [[i > j for i in range(len(corr_matrix.columns))] for j in range(len(corr_matrix.columns))],
            columns=corr_matrix.columns,
            index=corr_matrix.index
        )

        # 提取下三角矩阵（排除对角线）
        """
        where()是pandas中的条件筛选方法
        保留掩码为True位置的原始值
        将掩码为False位置的值替换为NaN
        """
        lower_triangle = corr_matrix.where(lower_triangle_mask)

        # 找到相关性大于0.8的特征对
        high_corr_pairs = []
        high_corr_features = set()

        for column in lower_triangle.columns:
            for row in lower_triangle.index:
                if lower_triangle.loc[row, column] > 0.8 and row != column:
                    high_corr_pairs.append((row, column, lower_triangle.loc[row, column]))
                    high_corr_features.add(row)
                    high_corr_features.add(column)

        self.logger.info(f"发现{len(high_corr_pairs)}对高相关性特征(相关系数>0.8)")

        # 记录高相关性特征对的详细信息
        for row, col, corr_value in high_corr_pairs:
            # 根据特征重要性判断哪个特征更重要
            if self.feature_importance is not None:
                imp_row = self.feature_importance[row] if row in self.feature_importance.index else 0
                imp_col = self.feature_importance[col] if col in self.feature_importance.index else 0

                if imp_row > imp_col:
                    self.logger.info(
                        f"高相关性特征对: '{row}' 和 '{col}' (相关性: {corr_value:.3f}), '{row}' 特征重要性更高")
                else:
                    self.logger.info(
                        f"高相关性特征对: '{row}' 和 '{col}' (相关性: {corr_value:.3f}), '{col}' 特征重要性更高")
            else:
                self.logger.info(f"高相关性特征对: '{row}' 和 '{col}' (相关性: {corr_value:.3f})")

        self.logger.info("特征相关性分析完成".center(50, "="))

        # 使用 seaborn 绘制 热力图 可视化 相关性矩阵
        self.logger.info("绘制相关性热力图")
        plt.figure(figsize=(12, 10))
        final_features_data = self.x[self.selected_features] if self.selected_features else self.x
        final_corr_matrix = final_features_data.corr()

        # 创建正确的 下三角 掩码
        mask = pd.DataFrame(
            [[i > j for i in range(len(final_corr_matrix.columns))] for j in range(len(final_corr_matrix.columns))],
            columns=final_corr_matrix.columns,
            index=final_corr_matrix.index
        )

        # 使用seaborn绘制热力图
        sns.heatmap(final_corr_matrix,
                    mask=mask,
                    annot=False,  # 不显示具体数值，避免图表过于拥挤
                    cmap='coolwarm',
                    center=0,
                    square=True,
                    fmt='.2f',
                    cbar_kws={"shrink": .8})

        plt.title('特征相关性热力图')
        plt.xticks(rotation=45, ha='right')
        plt.yticks(rotation=0)
        plt.tight_layout()
        plt.savefig('../data/fig/correlation_heatmap.png', dpi=300, bbox_inches='tight')
        plt.show()

        return high_corr_pairs


if __name__ == '__main__':
    fe = FeatureEngineering('../data/train.csv')
    fe.preprocess()
    fe.feature_importance_rf()
    fe.select_features_rfe()
    fe.features_corr()
